import os
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET
from glob import glob
from scipy.ndimage import gaussian_filter


def generate_density_map(shape, label, sigma = 8):
    """
    Generate the density map based on the dot annotations provided by the label.
    """
    height, width = shape
    label = torch.tensor(label).float()
    
    density_map = torch.zeros((1, height, width), dtype=torch.float32)

    if len(label) > 0:
        assert len(label.shape) == 2 and label.shape[1] == 2, f"label should be a Nx2 tensor, got {label.shape}."
        label_ = label.long()
        label_[:, 0] = label_[:, 0].clamp(min=0, max=width - 1)
        label_[:, 1] = label_[:, 1].clamp(min=0, max=height - 1)
        density_map[0, label_[:, 1], label_[:, 0]] = 1.0

    if sigma is not None:
        assert sigma > 0, f"sigma should be positive if not None, got {sigma}."
        density_map = torch.from_numpy(gaussian_filter(density_map, sigma=sigma))

    return density_map.squeeze()




savepath = "./01_draw_backbone_img"
if not os.path.exists(savepath):
    os.makedirs(savepath)

image = '/data/store1/nzd/tir_cc/datasets/DroneRGBT/Test/Infrared/1125R.jpg'
img = plt.imread(image)
shape = img.shape[:2]

crop_area = {
    'x_start': 180,   # 左边界
    'x_end': 280,     # 右边界
    'y_start': 170,   # 上边界
    'y_end': 270      # 下边界
}

points = [[238, 190], [212, 231], [252, 231]]



fig, ax = plt.subplots()
ax.imshow(img)

# 设置坐标轴显示范围
ax.set_xlim(crop_area['x_start'], crop_area['x_end'])
ax.set_ylim(crop_area['y_end'], crop_area['y_start'])  # 注意Y轴方向

# 隐藏坐标轴
ax.axis('off')

# 保存图像（去除白边，保持原始尺寸）
plt.savefig(f"{savepath}/thermal_img", bbox_inches='tight', pad_inches=0, dpi=300)
plt.show()
# 关闭图形释放内存
plt.close()



fig, ax = plt.subplots()
ax.imshow(img)

for (x, y) in points:
    ax.plot(x, y, 'ro', markersize=5)  # 'ro'表示红色圆圈

# 设置坐标轴显示范围
ax.set_xlim(crop_area['x_start'], crop_area['x_end'])
ax.set_ylim(crop_area['y_end'], crop_area['y_start'])  # 注意Y轴方向

# 隐藏坐标轴
ax.axis('off')

# 保存图像（去除白边，保持原始尺寸）
plt.savefig(f"{savepath}/point_img", bbox_inches='tight', pad_inches=0, dpi=300)
plt.show()
# 关闭图形释放内存
plt.close()



import numpy as np

color = [0.0, 0.0, 1.0]  # 纯蓝色 (RGB归一化值)
blue_background = np.ones((shape[0], shape[1], 3)) * color  # 创建三维数组

fig, ax = plt.subplots()
ax.imshow(blue_background, alpha=0.7)

for (x, y) in points:
    ax.plot(x, y, 'yo', markersize=5)  # 'ro'表示红色圆圈

# 设置坐标轴显示范围
ax.set_xlim(crop_area['x_start'], crop_area['x_end'])
ax.set_ylim(crop_area['y_end'], crop_area['y_start'])  # 注意Y轴方向

# 隐藏坐标轴
ax.axis('off')

# 保存图像（去除白边，保持原始尺寸）
plt.savefig(f"{savepath}/point_excep", bbox_inches='tight', pad_inches=0, dpi=300)
plt.show()
# 关闭图形释放内存
plt.close()



density = generate_density_map(shape, points, sigma=8)

fig, ax = plt.subplots()
ax.imshow(density)

# 设置坐标轴显示范围
ax.set_xlim(crop_area['x_start'], crop_area['x_end'])
ax.set_ylim(crop_area['y_end'], crop_area['y_start'])  # 注意Y轴方向

# 隐藏坐标轴
ax.axis('off')

# 保存图像（去除白边，保持原始尺寸）
plt.savefig(f"{savepath}/density", bbox_inches='tight', pad_inches=0, dpi=300)
plt.show()
# 关闭图形释放内存
plt.close()